from pathlib import Path
from typing import Union, Literal, List, Dict

import orjson

from .decorator import param2path
from .img import ImageTools
from .utils import P_TYPE


class Point(object):
    """
    label me 坐标点
    """

    def __init__(self, x: float, y: float):
        self.x = 0 if x < 0 else x
        self.y = 0 if y < 0 else y

    def to_list(self, x_max: float, y_max: float) -> List:
        """
        数据转列表 并限制坐标最大值
        Args:
            x_max:
            y_max:

        Returns:

        """
        _x, _y = self.x, self.y
        return [
            _x if _x < x_max else x_max - 1,
            _y if _y < y_max else y_max - 1,
        ]


class Shape(object):
    """
    label me 标注框形状
    """
    SHAPE_TYPE = Literal['rectangle', 'polygons']

    def __init__(self, label: str, points: List[Union[Point, List]],
                 group_id: str = None, shape_type: SHAPE_TYPE = 'rectangle',
                 flags: dict = None):
        if flags is None:
            flags = {}

        self.label = label
        self.points = points
        for _i, _point in enumerate(self.points):
            if isinstance(_point, list):
                self.points[_i] = Point(*_point)
        self.group_id = group_id
        self.shape_type = shape_type
        self.flags = flags

    def to_dict(self, **kwargs) -> Dict:
        """
        数据转字典
        Args:
            **kwargs:

        Returns:

        """
        return {
            'label': self.label,
            'points': [_point.to_list(**kwargs) for _point in self.points],
            'group_id': self.group_id,
            'shape_type': self.shape_type,
            'flags': self.flags
        }


class LabelMe(object):
    """
    label me 数据结构
    """

    @param2path({'p': 1, 'image_path': -1})
    def __init__(self, p: Union[str, Path] = None, version: str = '4.5.7',
                 flags: dict = None, shapes: List[Union[Shape, Dict]] = None,
                 image_data: str = None, image_width: int = None,
                 image_height: int = None, image_path: P_TYPE = None):
        if flags is None:
            flags = {}
        if shapes is None:
            shapes = []

        self.version = version
        self.flags = flags
        self.shapes = shapes
        for _i, _shape in enumerate(self.shapes):
            if isinstance(_shape, dict):
                self.shapes[_i] = Shape(**_shape)

        if p is not None:
            self.path_c = p
            self.image_tools = ImageTools(self.path_c)
            self.image_base64 = self.image_tools.to_base64
            self.width, self.height = self.image_tools.size
        else:
            self.path_c = image_path
            self.image_base64 = image_data
            self.width, self.height = image_width, image_height

    @property
    def to_dict(self) -> Dict:
        """
        数据转字典
        Returns:

        """
        _w, _h = self.width, self.height
        return {
            'version': self.version,
            'flags': self.flags,
            'shapes': [
                _shape.to_dict(x_max=_w, y_max=_h) for _shape in self.shapes],
            'imagePath': self.path_c.name.lower(),
            'imageData': self.image_base64,
            'imageHeight': _h,
            'imageWidth': _w,
        }

    @classmethod
    @param2path({'p': 1})
    def from_json(cls, p: Union[str, Path]) -> 'LabelMe':
        """
        获取标注文件内容
        Args:
            p: 文件地址

        Returns:

        """
        _json = orjson.loads(p.read_text(encoding='utf-8'))
        return cls(**{
            'version': _json['version'],
            'flags': _json['flags'],
            'shapes': _json['shapes'],
            'image_path': _json['imagePath'],
            'image_data': _json['imageData'],
            'image_height': _json['imageHeight'],
            'image_width': _json['imageWidth'],
        })

    @param2path({'p': 1}, ['p'])
    def save(self, p: P_TYPE) -> 'LabelMe':
        """
        文件保存
        Args:
            p: 保存文件夹路径

        Returns:

        """
        p.joinpath(f"{self.path_c.stem.lower()}.json").write_bytes(
            orjson.dumps(self.to_dict,
                         option=(orjson.OPT_APPEND_NEWLINE
                                 | orjson.OPT_INDENT_2)
                         )
        )
        return self
